#include <stdafx.h>
#include "DriverLoad.h"
#include <winioctl.h>
#pragma comment(lib,"Advapi32.lib")
#define DRIVERNAME _T("FirstDriver.sys")

CDriverLoad::CDriverLoad(void)
{
	m_hSCManager = NULL;
	m_hService = NULL;
	m_hDriver = INVALID_HANDLE_VALUE;
	m_IsStarted = FALSE;
	m_bOpened = FALSE;
}

CDriverLoad::~CDriverLoad(void)
{
	Remove();
	if (m_hDriver)
	{
		CloseHandle(m_hDriver);
		m_hDriver = INVALID_HANDLE_VALUE;
	}
	if (m_hService)
	{
		CloseServiceHandle(m_hService);
		m_hService = NULL;
	}
	if (m_hSCManager)
	{
		CloseServiceHandle(m_hSCManager);
		m_hSCManager = NULL;
	}
}

BOOL CDriverLoad::UnInstall(PTCHAR pServiceName)
{
	SERVICE_STATUS ss;
	SC_HANDLE hService = NULL;
	SC_HANDLE hSCManager = NULL;

	hSCManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
	if (NULL == hSCManager)
	{
		goto RET;
	}

	hService = OpenService(hSCManager, pServiceName, SERVICE_ALL_ACCESS);
	if (NULL == hService)
	{
		goto RET;
	}
	if (!ControlService(hService, SERVICE_CONTROL_STOP, &ss))
	{
		//	goto RET;
	}
	if (!DeleteService(hService))
	{
		//	goto RET;
	}
	TCHAR			SysDrivPath[MAX_PATH] = { 0 };
	
	GetSystemDirectory(SysDrivPath, MAX_PATH);
	_tcscat_s(SysDrivPath, _tcslen(SysDrivPath) + _tcslen(_T("\\drivers\\")) + 1, _T("\\drivers\\"));
	_tcscat_s(SysDrivPath, _tcslen(SysDrivPath) + _tcslen(DRIVERNAME) + 1, DRIVERNAME);
	DeleteFile(SysDrivPath);
RET:
	if (hService)
	{
		CloseServiceHandle(hService);
	}
	if (hSCManager)
	{
		CloseServiceHandle(hSCManager);
	}
	return TRUE;
}

BOOL CDriverLoad::Install(PTCHAR PInfPath, PTCHAR pServiceName)
{
	TCHAR			SysDrivPath[MAX_PATH] = { 0 };
	WIN32_FIND_DATA FindFileData;
	int				count = 100;
	UnInstall(pServiceName);
	ShellExecute(NULL, _T("install"), PInfPath, NULL, NULL, SW_SHOW);
	GetSystemDirectory(SysDrivPath, MAX_PATH);
	_tcscat_s(SysDrivPath, _tcslen(SysDrivPath) + _tcslen(_T("\\drivers\\")) + 1, _T("\\drivers\\"));
	_tcscat_s(SysDrivPath, _tcslen(SysDrivPath) + _tcslen(DRIVERNAME) + 1, DRIVERNAME);

	while (FindFirstFile(SysDrivPath, &FindFileData) == INVALID_HANDLE_VALUE && count > 0)
	{
		count--;
		Sleep(10);
	}
	Sleep(300);
	m_hSCManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
	if (NULL == m_hSCManager)
	{
		m_dwLastError = GetLastError();
		return FALSE;
	}

	m_hService = OpenService(m_hSCManager, pServiceName, SERVICE_ALL_ACCESS);
	if (NULL == m_hService)
	{
		CloseServiceHandle(m_hSCManager);
		m_hSCManager = NULL;
		return FALSE;
	}

	return TRUE;
}
BOOL CDriverLoad::Start()
{
	if (!m_hService)
	{
		return FALSE;
	}
	if (!StartService(m_hService, NULL, NULL))
	{
		m_dwLastError = GetLastError();
		return FALSE;
	}
	m_IsStarted = TRUE;
	return TRUE;
}

BOOL CDriverLoad::IsRemoved()
{
	if (m_hSCManager || m_hService)
	{
		return FALSE;
	}
	else
	{
		return TRUE;
	}
}

BOOL CDriverLoad::IsInstalled()
{
	if (m_hSCManager && m_hService)
	{
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

BOOL CDriverLoad::Stop()
{
	SERVICE_STATUS ss;
	if (!m_IsStarted)
	{
		return TRUE;
	}
	m_bOpened = FALSE;
	if (!ControlService(m_hService, SERVICE_CONTROL_STOP, &ss))
	{
		m_dwLastError = GetLastError();
		return FALSE;
	}
	m_IsStarted = FALSE;
	return TRUE;

}

BOOL CDriverLoad::Remove()
{
	close();
	Stop();
	if (!DeleteService(m_hService))
	{
		m_dwLastError = GetLastError();
		return FALSE;
	}
	if (m_hDriver != INVALID_HANDLE_VALUE)
	{
		CloseHandle(m_hDriver);
		m_hDriver = INVALID_HANDLE_VALUE;
	}
	if (m_hService)
	{
		CloseServiceHandle(m_hService);
		m_hService = NULL;
	}
	if (m_hSCManager)
	{
		CloseServiceHandle(m_hSCManager);
		m_hSCManager = NULL;
	}
	return TRUE;
}

BOOL CDriverLoad::Open(PTCHAR pLinkName)
{
	if (m_hDriver != INVALID_HANDLE_VALUE)
		return TRUE;
	m_hDriver = CreateFile(pLinkName, GENERIC_READ | GENERIC_WRITE, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
	if (m_hDriver != INVALID_HANDLE_VALUE)
	{
		m_bOpened = TRUE;
		return TRUE;
	}
	else
	{
		return FALSE;
	}

}

BOOL CDriverLoad::close()
{
	if (m_hDriver != INVALID_HANDLE_VALUE)
	{
		if (CloseHandle(m_hDriver))
		{
			m_hDriver = INVALID_HANDLE_VALUE;
			return TRUE;
		}
	}
	return FALSE;
}

BOOL CDriverLoad::IoControl(DWORD dwIoCode, PVOID InBuff, DWORD InBuffLen, PVOID OutBuff, DWORD OutBuffLen, DWORD *RealRetBytes)
{
	DWORD dw;
	if (!m_IsStarted)
	{
		return FALSE;
	}
	BOOL b = DeviceIoControl(m_hDriver, dwIoCode, InBuff, InBuffLen, OutBuff, OutBuffLen, &dw, NULL);
	if (RealRetBytes)
		*RealRetBytes = dw;
	return b;
}

